import os

import paddlex as pdx

from src.config.config import OUTPUT_DIR, DATA_ROOT

if __name__ == '__main__':
    model = pdx.load_model(os.path.join(OUTPUT_DIR, 'yolov3_darknet53', 'best_model'))
    base_path = os.path.join(DATA_ROOT, 'insect_det', 'JPEGImages')
    for name in os.listdir(base_path):
        image_name = os.path.join(base_path, name)
        result = model.predict(image_name)
        pdx.det.visualize(image_name, result, threshold=0.5,
                          save_dir=os.path.join(OUTPUT_DIR, 'yolov3_darknet53', 'results'))
